#include <stdio.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>
#include <sys/time.h>
#include <time.h>
#include <mpi.h>

#define EPSILON 1.0e-5
#define REAL_T float
#define C 0.2


#define TAG_FIRST_LINE 42
#define TAG_LAST_LINE 43
#define TAG_LEFT_LINE 44
#define TAG_RIGHT_LINE 45

/* Pour obtenir le temps "horloge" : */
double my_gettimeofday(){
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec + (tv.tv_usec * 1.0e-6L);
}

static inline REAL_T calcul_iter (REAL_T **grid[2], int grid_nb, int i, int j, int h, int w, int first_index, int first_jndex) {
        return ((1.0 - 4.0 * C) * grid[grid_nb][i][j]
          + C * (j>0   ? grid[grid_nb][i][j-1] : grid[grid_nb][i][j])
          + C * (j+first_jndex<w-1 ? grid[grid_nb][i][j+1] : grid[grid_nb][i][j])
          + C * (i>0   ? grid[grid_nb][i-1][j] : grid[grid_nb][i][j]) 
          + C * (i+first_index<h-1 ? grid[grid_nb][i+1][j] : grid[grid_nb][i][j]));
}

int main(int argc, char **argv){

    /******************/
    /* INITIALISATION */
    /******************/
    
    int n = 0;
    int w = 0;
    int h = 0;
    int i = 0, j = 0;
    int step = 0;

    int grid_nb = 0, inv_grid_nb=0;
    REAL_T **grid[2] = {NULL, NULL};

    REAL_T max_error = 0.0;

    double t_debut = 0.0, t_fin = 0.0;
    double t_total_debut = 0.0, t_total_fin = 0.0;

    /* Variables MPI */
    int rank;
    int size;
    int h_local, first_index, w_local, first_jndex, w_size;
    int iter;
    REAL_T max_error_local = 0.0;
    int h_debut, h_fin, w_debut, w_fin;
    MPI_Status status;
    int dims[2];
    int periods[2];
    int coords[2];
    int reorder;
    MPI_Comm cart_comm;
    REAL_T * buff_grid1, *buff_grid2;
    int buff_size;
    int rank_left,rank_up, rank_down, rank_right;

    MPI_Request req_cr_first, req_cr_last, req_cs_first, req_cs_last;
    MPI_Request req_lr_first, req_lr_last, req_ls_first, req_ls_last;

    /* Valeurs par defaut : */
    w = 240;
    h = 240;
    dims[0] = 3; // 3 per lines
    dims[1] = 3; // 3 per cols
    periods[0] = 1;
    periods[1] = 1;
    reorder = 0;

    /* Parametres : */
    if( argc > 1) h             = atoi(argv[1]);
    if( argc > 2) w             = atoi(argv[2]);
    if( argc > 3) dims[0]       = atoi(argv[3]);
    if( argc > 4) dims[1]       = atoi(argv[4]);
    if( argc > 5) reorder       = atoi(argv[5]);

    /* Initialisation MPI */
    MPI_Init(&argc, &argv);
    
    /* Create Cartesian topology */
    MPI_Cart_create(MPI_COMM_WORLD, 2, dims, periods, reorder, &cart_comm);
    
    MPI_Comm_size(cart_comm, &size);
    MPI_Comm_rank(cart_comm, &rank);
    MPI_Cart_coords(cart_comm, rank, 2, coords);    
    
    if (size != dims[0]*dims[1])
    {
        printf("Please run with %d processes.\n",(dims[0]*dims[1]));fflush(stdout);
        MPI_Abort(MPI_COMM_WORLD, 1);
    }
    
    int c_tmp[2];
    c_tmp[0] = coords[0];c_tmp[1] = coords[1]-1; // {+0, -1}
    MPI_Cart_rank(cart_comm, c_tmp, &rank_up); 
    c_tmp[0] = coords[0];c_tmp[1] = coords[1]+1; // {+0, +1}
    MPI_Cart_rank(cart_comm, c_tmp, &rank_down); 
    c_tmp[0] = coords[0]-1;c_tmp[1] = coords[1]; // {-1, +0}
    MPI_Cart_rank(cart_comm, c_tmp, &rank_left); 
    c_tmp[0] = coords[0]+1;c_tmp[1] = coords[1]; // {+1, +0}
    MPI_Cart_rank(cart_comm, c_tmp, &rank_right); 
     
    /* Calcul selon h */
    first_index = ((h/dims[1]) * coords[1]) - (coords[1]>0);
    h_local = (h/dims[1]) + (coords[1]>0) + (coords[1] < (dims[1]-1));
    
    h_debut = (coords[1]>0);
    h_fin = h_local - (coords[1] < (dims[1]-1));
    
    /* Calcul selon w */
    first_jndex = ((w/dims[0]) * coords[0]) - (coords[0]>0);
    w_local = (w/dims[0]) + (coords[0]>0) + (coords[0]<dims[0]-1);
    
    w_debut = (coords[0]>0);
    w_fin = w_local - (coords[0] < (dims[0]-1));
    
    w_size = w/dims[0];
    
    /* buffer for a col : h dims*/
    buff_size = h/dims[1];
    
    if(rank == 0) {
        /* Demarrage horloge : */
        t_total_debut = my_gettimeofday();

        /* Affichage informations diverses : */
        //printf("Dims : (%d, %d) \t Nb points : %i \t Taille memoire : %f Mo X2\n",w, h, w*h, (float) w*h*sizeof(REAL_T) / (1024*1024));
    }
    
    //printf("%d(%d,%d) - lignes %d à %d, cols %d à %d - up:%d, down:%d, left:%d; right:%d \n", rank, coords[0], coords[1], h_debut+first_index, h_fin+first_index, w_debut+first_jndex, w_fin+first_jndex, rank_up, rank_down, rank_left, rank_right);
    
    /* Allocation memoire : h_local et w_local */
    for (grid_nb=0; grid_nb<=1; grid_nb++){
        grid[grid_nb] = (REAL_T **) malloc(h_local*sizeof(REAL_T *));
        if (grid[grid_nb] == NULL){
            fprintf(stderr, "Echec allocation memoire pour 'grid[%d]'.\n", grid_nb);
        }
        for (i=0; i<h_local; i++){
            grid[grid_nb][i] = (REAL_T *) malloc(w_local*sizeof(REAL_T));
            if (grid[grid_nb][i] == NULL){
                fprintf(stderr, "Echec allocation memoire pour 'grid[%d][%d]'.\n", grid_nb, i);
            }
            /* set to 0: */
            memset(grid[grid_nb][i], 0, w_local*sizeof(REAL_T));
        } /* for i */
    } /* for grid_nb */

    /* Allocation buffer d'envoi des colonnes */
    buff_grid1 = (REAL_T *) malloc(buff_size*sizeof(REAL_T));
    
    /* Allocation buffer de reception des colonnes */
    buff_grid2 = (REAL_T *) malloc(buff_size*sizeof(REAL_T));

    /* Initialisation de la grille : */
    grid_nb=0;
    for (i=0; i<h_local; i++){
        for (j=0; j<w_local; j++){
            if((i+first_index)%100 == 50)
                grid[grid_nb][i][j] = (REAL_T) 100.0;
            else
                grid[grid_nb][i][j] = (REAL_T) 25.0;
        }
    }

    /* Calcul */

    iter = 0;
    
    do {
        inv_grid_nb = 1-grid_nb;

        if(rank == 0) {
            /* Debut horloge pour cette iteration : */
            t_debut = my_gettimeofday();
        }

        /* mise a zero de grid[inv_grid_nb]: */
        for (i=0; i<h_local; i++){
            memset(grid[inv_grid_nb][i], 0, w_local*sizeof(REAL_T));    
        } /* for i */
	
	 /* Receive de la ligne 0 de (rank_up) -> LAST_LINE */
        if(coords[1]>0) {
	  MPI_Irecv(&grid[inv_grid_nb][0][w_debut], w_size, MPI_FLOAT, rank_up, TAG_LAST_LINE, cart_comm, &req_lr_last);
        }

        /* Receive de la ligne h_local-1 de (rank_down) -> FIRST_LINE */
        if(coords[1] <(dims[1]-1)) {
	  MPI_Irecv(&grid[inv_grid_nb][h_local-1][w_debut], w_size, MPI_FLOAT, rank_down, TAG_FIRST_LINE, cart_comm, &req_lr_first);
        }

	 /* Receive de la col 0 de (rank_left) -> RIGHT_LINE */
        if(coords[0]>0) {
	  MPI_Irecv(buff_grid1, buff_size, MPI_FLOAT, rank_left, TAG_RIGHT_LINE, cart_comm, &req_cr_first);
	  /* Dé-Bufferisation */
	  for(i=0;i<buff_size;i++) {
	    grid[inv_grid_nb][i+h_debut][0] = buff_grid1[i];
	  }
        }

	/* Receive de la col w_local-1 de (rank_right) -> LEFT_LINE */
        if(coords[0] <(dims[0]-1)) {
	  MPI_Irecv(buff_grid2, buff_size, MPI_FLOAT, rank_right, TAG_LEFT_LINE, cart_comm, &req_cr_last);
	  /* Dé-Bufferisation */
	  for(i=0;i<buff_size;i++) {
	    grid[inv_grid_nb][i+h_debut][w_local-1] = buff_grid2[i];
	  }
        }

	/* Calcul de la 1ère ligne */
	for(j=w_debut; j<w_fin; j++){
	  i = h_debut;
	  grid[inv_grid_nb][i][j] = calcul_iter (grid, grid_nb, i, j, h, w, first_index, first_jndex); 
	}
	
        /* Envoi de la ligne 1 à (rank_up) -> FIRST_LINE */
        if(coords[1]>0) {
	  MPI_Isend(&grid[inv_grid_nb][1][w_debut], w_size, MPI_FLOAT, rank_up, TAG_FIRST_LINE, cart_comm, &req_ls_first);
        }

	/* Calcul de la derniere ligne */
	for(j=w_debut; j<w_fin; j++){
	  i = h_fin-1;
	  grid[inv_grid_nb][i][j] = calcul_iter (grid, grid_nb, i, j, h, w, first_index, first_jndex); 
	}
	
        /* Envoi de la ligne h_local-2 à (rank_down) -> LAST_LINE */
        if(coords[1] <(dims[1]-1)) {
	  MPI_Isend(&grid[inv_grid_nb][h_local-2][w_debut], w_size, MPI_FLOAT, rank_down, TAG_LAST_LINE, cart_comm, &req_ls_last);
        }
	
	/* Calcul de la 1ère colonne */
	for(i=h_debut; i<h_fin; i++){
	  j = w_debut;
	  grid[inv_grid_nb][i][j] = calcul_iter(grid, grid_nb, i, j, h, w, first_index, first_jndex); 
	}
	
        /* Envoi de la col 1 à (rank_left) -> LEFT_LINE */
        if(coords[0]>0) {
	  /* Bufferisation */
	  for(i=0;i<buff_size;i++) {
	    buff_grid1[i] = grid[inv_grid_nb][i+h_debut][1];
	  }
	  
	  MPI_Isend(buff_grid1, buff_size, MPI_FLOAT, rank_left, TAG_LEFT_LINE, cart_comm, &req_cs_first);
        }

	/* Calcul de la dernière colonne */
	for(i=h_debut; i<h_fin; i++){
	  j = w_fin-1;
	  grid[inv_grid_nb][i][j] = calcul_iter(grid, grid_nb, i, j, h, w, first_index, first_jndex); 
	}
	
        /* Envoi de la col w_fin-2 à (rank_right) -> RIGHT_LINE */
        if(coords[0] <(dims[0]-1)) {
	  /* Bufferisation */
	  for(i=0;i<buff_size;i++) {
	    buff_grid1[i] = grid[inv_grid_nb][i+h_debut][w_fin-2];
	  }
	  
	  MPI_Isend(buff_grid1, buff_size, MPI_FLOAT, rank_right, TAG_RIGHT_LINE, cart_comm, &req_cs_last);
        }
        
	/* Calcul reste */
	for (i=h_debut+1; i<h_fin-1; i++){
	  for (j=w_debut+1; j<w_fin-1; j++){
	    grid[inv_grid_nb][i][j] = calcul_iter(grid, grid_nb, i, j, h, w, first_index, first_jndex); 
	  } /* for j */
        } /* for i */

        /* calcul de l'erreur */
        max_error_local = 0.0;
        for (i=h_debut; i<h_fin; i++){
            for (j=w_debut; j<w_fin; j++){
                REAL_T error = fabs(grid[grid_nb][i][j] - grid[inv_grid_nb][i][j]);
                if (error > max_error_local){ 
                    max_error_local = error; 
                }
            } /* for j */
        } /* for i */
	
        /* Réduction sur le max des max_error_local */
            // Attention, synchronisation obligatoire sans recouvrement possible 
            // A voir si autre possibilité non systématique envisageable 
        MPI_Allreduce(&max_error_local, &max_error, 1, MPI_FLOAT, MPI_MAX, cart_comm);
        
        if(rank == 0) {
            /* Arret horloge pour cette iteration : */
            t_fin = my_gettimeofday();

            /* Affichage pour cette iteration : */
            //printf("Iteration %d : \t delta = %.3e \t temps = %.1f s\n",++n, max_error, t_fin - t_debut);
        }

	// Attendre les opérations non bloquante
	if(coords[1]>0){
	  MPI_Wait(&req_lr_last, &status);
	  MPI_Wait(&req_ls_first, &status);
	}
	if(coords[1] <(dims[1]-1)){
	  MPI_Wait(&req_lr_first, &status);
	  MPI_Wait(&req_ls_last, &status);
	}
	
	if(coords[0] > 0){
	  MPI_Wait(&req_cr_first, &status);
	  MPI_Wait(&req_cs_first, &status);
	}
	
	if(coords[0] < (dims[0]-1)){
	  MPI_Wait(&req_cr_last, &status);
	  MPI_Wait(&req_cs_last, &status);
	}
        
        /* update 'grid_nb': */
        grid_nb = inv_grid_nb;
        
        /* itération suivante */
        iter++;
    } while (max_error > EPSILON && iter < 40);
    /***** Fin du calcul *****/

    if(coords[0] == 0){
      for(i=h_debut;i<h_fin;i++) {
	printf("%d %3.5f\n", coords[1], grid[inv_grid_nb][i][1]);
      }
    }
    
    if(rank == 0) {
      /* Arret horloge : */
      t_total_fin = my_gettimeofday();  
      
      /* Affichage des resultats : */
      /*
      printf("\nNombre d'iterations : %d\n", n); 
      printf("Erreur (delta) = %.3e\n", max_error); 
      printf("Temps total : %.1f s\n", t_total_fin - t_total_debut); 
      */
    }

    /* Liberation memoire : */
    for (grid_nb=0; grid_nb<=1; grid_nb++){
      for (i=0; i<h_local; i++){
	free(grid[grid_nb][i]);
      }
      
      free(grid[grid_nb]);
    }
    
    free(buff_grid1);
    
    MPI_Finalize();

    return EXIT_SUCCESS;
}
